package nl.cypherpunk.statelearner.fuzzing;

import net.automatalib.automata.transducers.MealyMachine;
import net.automatalib.words.Alphabet;
import net.automatalib.words.Word;
import nl.cypherpunk.statelearner.StateLearnerSUL;
import nl.cypherpunk.statelearner.mealyDot.MealyDot;
import nl.cypherpunk.statelearner.mealyDot.MealyState;
import nl.cypherpunk.statelearner.mealyDot.OutputTransition;

import java.io.*;
import java.nio.charset.Charset;
import java.text.SimpleDateFormat;
import java.util.*;

public class FuzzingBasedTesting {
    MealyDot mealyAutoamta;
    MealyState initialState;
    Random rand = new Random();
    Alphabet<String> alphabet;
    Alphabet<String> meaning;
    TLSSUT sut;
    Map<String, LinkedHashMap< byte[],Double>> seed;
    ArrayList<List<byte[]>> standardValue;
    ArrayList<Word<String>> standardKey;
    ArrayList<Boolean> standardStatus;
    ArrayList<Double> standardScore;
    PrintStream log;
    String server_output;
    String seed_dir;
    String coverage_cmd;
    int initCoverage = 0;
    int total_num = 0;

    public FuzzingBasedTesting(String dir, String dotFileName, TLSSUT sut, Alphabet<String> inputAlphabet, String cmd) throws Exception {
        mealyAutoamta = new MealyDot(dotFileName);
        initialState = mealyAutoamta.getInitialState();
        this.sut = sut;
        this.alphabet = inputAlphabet;
        this.meaning = sut.getMeaning();
        this.server_output = dir + File.separator + "server.log";
        File serverLog = new File(this.server_output);
        this.seed_dir = dir + File.separator + "Seed";
        File outputDir = new File(seed_dir);
        outputDir.mkdirs();
        this.coverage_cmd = cmd;
        standardValue = new ArrayList<>();
        standardKey = new ArrayList<>();
        standardStatus = new ArrayList<>();
        standardScore = new ArrayList<>();
        seed = new HashMap<>();
        for (int i = 0; i < alphabet.size(); i++) {
            seed.put(alphabet.getSymbol(i), new LinkedHashMap<>());
        }
        seed.put("HeartbeatRequest", new LinkedHashMap<>());
        File outputFile = new File(dir, "fuzzing.log");
        log = new PrintStream(outputFile);
    }

    public void randomWalkWithFuzzing(int number, int time) throws Exception,IOException, ClassNotFoundException, InterruptedException {
        //Traverse all states based on depth-first strategy, and record the concrete messages of the set Q
        initStandard();
        //initCoverage = 0;
        int prevCoverage = initCoverage;
        this.sut.setInitScore(initCoverage*1.0);
        int num = number;
        this.sut.pre();
        Date startDate = new Date();
        String timestamp = String.valueOf(startDate.getTime()/1000);
        int start =  Integer.valueOf(timestamp);
        //select transition path
        while(true){
            // random select state chain
            this.sut.pre();
            int stateIndex = getRandomState(standardScore);
            Word<String> inputList = standardKey.get(stateIndex);
            List<byte[]> messageList = standardValue.get(stateIndex);
            ArrayList<byte[]> testCase = new ArrayList<>();
            MealyState currState = initialState;
            //select the state to mutate
            int mutateState = rand.nextInt(inputList.length());
            int k = 0;
            for ( ; k < mutateState; k++) {
                String output = sut.stepConcrete(inputList.getSymbol(k), seed.get(inputList.getSymbol(k)), false);
                byte[] tempSeed = sut.tls.currentSeed.get(inputList.getSymbol(k));
                if(tempSeed != null )
                    testCase.add(tempSeed);
                OutputTransition outputTransition = currState.getOutputTransition(inputList.getSymbol(k));
                currState = outputTransition.state;
                log.print(inputList.getSymbol(k) + " ");
            }
            String mutateInput = null;
            String mutateOutput = null;
            byte[] saveSeed = null;
            if(mutateState%2 == 0){
                mutateInput = inputList.getSymbol(mutateState);
                mutateOutput = sut.stepConcrete(mutateInput, seed.get(mutateInput), true);
                saveSeed = sut.tls.currentSeed.get(inputList.getSymbol(mutateState));
            }
            else
            {
                mutateInput = "HeartbeatRequest";
                mutateOutput = sut.stepConcrete(mutateInput, seed.get(mutateInput), true);
                saveSeed = sut.tls.currentSeed.get(inputList.getSymbol(mutateState));
            }
            if(saveSeed != null )
                testCase.add(saveSeed);
            log.print(mutateInput + "(mutate) ");
            k = k + 1;
            for(; k < inputList.length(); k++ ){
                String output = sut.stepConcrete(inputList.getSymbol(k), seed.get(inputList.getSymbol(k)), false);
                byte[] tempSeed = sut.tls.currentSeed.get(inputList.getSymbol(k));
                if(tempSeed != null )
                    testCase.add(tempSeed);
                OutputTransition outputTransition = currState.getOutputTransition(inputList.getSymbol(k));
                currState = outputTransition.state;
                log.print(inputList.getSymbol(k) + " ");
            }
            log.println(Arrays.toString(saveSeed));
            log.println(mutateOutput);
            log.println("------------");
            total_num++;
            int coverage = this.sut.getCoverage();
            if(coverage == 1)
                seedSave(testCase, this.seed_dir + File.separator +  "crash-" + (Arrays.toString(saveSeed)).hashCode());
            if(saveSeed!=null)
            {
                if(coverage > prevCoverage)
                {
                    if(coverage==0)
                        coverage = initCoverage;
                    seed.get(mutateInput).put(saveSeed,Double.valueOf(coverage));
                    System.out.println(coverage);
                    if(coverage > prevCoverage)
                    {
                        seedSave(testCase, this.seed_dir + File.separator +String.format("%06d",total_num) +mutateOutput + "-" + (Arrays.toString(saveSeed)).hashCode()+"-" + String.valueOf(coverage-prevCoverage));
                        standardScore.set(stateIndex, standardScore.get(stateIndex)+1);
                    }
                    prevCoverage = coverage;
                }
            }
            Date endDate = new Date();
            String endstamp = String.valueOf(endDate.getTime()/1000);
            int end =  Integer.valueOf(endstamp);
            if(end - start > time)
                break;
        }

        /*
        for(int j = 0; j < standardKey.size(); j++)
        {
            int stateIndex = j;
            System.out.println("The round of Fuzzing : "+ j);
            Word<String> inputList = standardKey.get(stateIndex);
            List<byte[]> messageList = standardValue.get(stateIndex);
            Map<Integer, LinkedHashMap<byte[], Double>> tmpSeed = new HashMap<>();
            for (int k = 0; k<inputList.length(); k++) {
                LinkedHashMap<byte[], Double> tmp = new LinkedHashMap<>();
                tmp.put(messageList.get(k), 10.0);
                tmpSeed.put(k,tmp);
            }
            //two stage: the first stage is mutate the last input, the second stage is mutate the next input

            for(int i = 0; i < num; i++) {
                this.sut.pre();
                ArrayList<byte[]> testCase = new ArrayList<>();
                MealyState currState = initialState;
                int k = 0;
                for (; k < inputList.length() - 1; k++) {
                    String output = sut.stepConcrete(inputList.getSymbol(k), tmpSeed.get(k), false);
                    byte[] tempSeed = sut.tls.currentSeed.get(inputList.getSymbol(k));
                    testCase.add(tempSeed);
                    OutputTransition outputTransition = currState.getOutputTransition(inputList.getSymbol(k));
                    currState = outputTransition.state;
                    log.print(inputList.getSymbol(k) + " ");
                }
                String input = inputList.getSymbol(k);
                String output = sut.stepConcrete(input, seed.get(input), true);
                byte[] saveSeed = sut.tls.currentSeed.get(inputList.getSymbol(k));
                testCase.add(saveSeed);
                log.println(inputList.getSymbol(k));
                //if not contains "closed",
                while (!output.contains("Closed")) {
                    int tmp_k = rand.nextInt(meaning.size());
                    String tmp = meaning.getSymbol(tmp_k);
                    log.print(tmp+" ");
                    output = sut.stepConcrete(tmp, seed.get(tmp), false);
                    byte[] tempSeed = sut.tls.currentSeed.get(meaning.getSymbol(tmp_k));
                    testCase.add(tempSeed);
                }
                log.println(Arrays.toString(saveSeed));
                log.println(output);
                log.println("------------");
                total_num++;
                int coverage = 0;
                String resultStr = new String();
                File file = new File(this.server_output);
                try {
                    BufferedReader bufferedReader = new BufferedReader(new FileReader(file));
                    String str = null;
                    while (null != (str = bufferedReader.readLine())) {
                        //the output of server exists many 0x00
                        String tmp = str.replaceAll("\\u0000", "");
                        if(tmp.contains("Coverage"))
                        {
                            try {
                                String cover = tmp.substring(9).trim();
                                coverage = Integer.parseInt(cover);

                            } catch (NumberFormatException e) {
                                e.printStackTrace();
                            }
                        }
                        resultStr = resultStr.concat(tmp + "\r\n");
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
                FileWriter fileWriter = new FileWriter(file, false);
                fileWriter.write("");
                fileWriter.flush();
                fileWriter.close();
                if(resultStr.contains("AddressSanitizer"))
                    seedSave(testCase, this.seed_dir + File.separator +  "crash-" + (Arrays.toString(saveSeed)).hashCode());
                if(saveSeed!=null)
                {
                    if(coverage > prevCoverage || i == 0)
                    {
                        if(coverage==0)
                            coverage = initCoverage;
                        seed.get(input).put(saveSeed,Double.valueOf(coverage));
                        System.out.println(coverage);
                        if(coverage > prevCoverage)
                            seedSave(testCase, this.seed_dir + File.separator +String.format("%06d",total_num) +output + "-" + (Arrays.toString(saveSeed)).hashCode()+"-" + String.valueOf(coverage-prevCoverage));
                        prevCoverage = coverage;
                    }
                }
            }

            if(!standardStatus.get(j))
            {
                for(int m = 0; m < meaning.size(); m++){
                    String input = meaning.getSymbol(m);
                    for(int i = 0; i < number; i++) {
                        // reach target state based on selected path
                        this.sut.pre();

                        String resultStr = new String();
                        ArrayList<byte[]> testCase = new ArrayList<>();
                        MealyState currState = initialState;
                        for (int k = 0; k < inputList.length(); k++) {
                            String output = sut.stepConcrete(inputList.getSymbol(k), tmpSeed.get(k), false);
                            byte[] tempSeed = sut.tls.currentSeed.get(inputList.getSymbol(k));
                            testCase.add(tempSeed);
                            OutputTransition outputTransition = currState.getOutputTransition(inputList.getSymbol(k));
                            currState = outputTransition.state;
                            log.print(inputList.getSymbol(k) + " ");
                        }
                        String output = sut.stepConcrete(input, seed.get(input), true);
                        byte[] saveSeed = sut.tls.currentSeed.get(input);
                        testCase.add(saveSeed);
                        log.println(input);
                        while (!output.contains("Closed"))
                        {
                            int tmp_k = rand.nextInt(meaning.size());
                            String tmp = meaning.getSymbol(tmp_k);
                            log.print(tmp+" ");
                            output = sut.stepConcrete(tmp, seed.get(tmp), false);
                            byte[] tempSeed = sut.tls.currentSeed.get(meaning.getSymbol(tmp_k));
                            testCase.add(tempSeed);
                        }
                        log.println(Arrays.toString(saveSeed));
                        log.println(output);
                        log.println("------------");
                        total_num++;

                        int coverage = 0;
                        File file = new File(this.server_output);
                        try {
                            BufferedReader bufferedReader = new BufferedReader(new FileReader(file));
                            String str = null;
                            while (null != (str = bufferedReader.readLine())) {
                                //the output of server exists many 0x00
                                String tmp = str.replaceAll("\\u0000", "");
                                if(tmp.contains("Coverage"))
                                {
                                    try {
                                        String cover = tmp.substring(9).trim();
                                        coverage = Integer.parseInt(cover);

                                    } catch (NumberFormatException e) {
                                        e.printStackTrace();
                                    }
                                }
                                resultStr = resultStr.concat(tmp + "\r\n");
                            }
                        } catch (IOException e) {
                            e.printStackTrace();
                        }
                        FileWriter fileWriter = new FileWriter(file, false);
                        fileWriter.write("");
                        fileWriter.flush();
                        fileWriter.close();
                        if (resultStr.contains("AddressSanitizer"))
                            seedSave(testCase, this.seed_dir + File.separator + "crash-" + (Arrays.toString(saveSeed)).hashCode());
                        if(saveSeed!=null)
                        {
                            if(coverage > prevCoverage || i == 0)
                            {
                                if(coverage==0)
                                    coverage = prevCoverage;

                                seed.get(input).put(saveSeed,Double.valueOf(coverage));
                                System.out.println(coverage);
                                if(coverage > prevCoverage)
                                    seedSave(testCase, this.seed_dir + File.separator +String.format("%06d",total_num) +output + "-" + (Arrays.toString(saveSeed)).hashCode()+"-" + String.valueOf(coverage-prevCoverage));
                                prevCoverage = coverage;
                            }
                        }
                    }
                }
            }
        }
        */
        log.println("The execs done : " + total_num);
        log.close();
        sut.exit();
    }

    private String randomStep() {
        return meaning.getSymbol(rand.nextInt(meaning.size()));
    }

    public void setInputAlphabet(Alphabet<String> inputAlphabet) {
        this.alphabet = inputAlphabet;
    }

    public void seedSave(List<byte[]> seed, String filename) {
        BufferedOutputStream bos = null;
        FileOutputStream fos = null;
        File file = null;
        try {
            file = new File(filename);
            if (file.exists()) {
                return;
            }
            fos = new FileOutputStream(file);
            bos = new BufferedOutputStream(fos);
            for (int i = 0;i < seed.size();i++)
            {
                byte[] tmp = intToByteArray(seed.get(i).length);
                bos.write(tmp,0,4);
                bos.write(seed.get(i), 0, seed.get(i).length);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (bos != null) {
                try {
                    bos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            if (fos != null) {
                try {
                    fos.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
        }
    }
    //bfs
    public void initStandard() throws IOException, ClassNotFoundException, InterruptedException
    {
        Queue<TransitionRecord> bfsQueue = new ArrayDeque<>();
        bfsQueue.add(new TransitionRecord(initialState,Word.epsilon(),Word.epsilon(), new LinkedList<byte[]>(),new LinkedList<MealyState>()));
        TransitionRecord curr = bfsQueue.poll();
        int j = 0;
        while(curr != null)
        {
            MealyState currState = curr.state;
            for(int i = 0; i < currState.getInputs().size(); i++)
            {
                this.sut.pre();
                ArrayList<byte[]> testCase = new ArrayList<>();
                List<byte[]> messageList = deepCopy(curr.messageList);
                List<MealyState> stateList = deepCopy(curr.stateList);
                String input = currState.getInputs().get(i);
                OutputTransition outputTransition = currState.getOutputTransition(input);
                MealyState targetState = outputTransition.state;
                if(targetState != null && !curr.stateList.contains(targetState)){
                    Iterator iter = curr.accessSequence.iterator();
                    while (iter.hasNext())
                    {
                        String tmpInput = iter.next().toString();
                        sut.stepConcrete(tmpInput, new LinkedHashMap<>(), false);
                    }
                    String output = sut.stepConcrete(input, new LinkedHashMap<>(), false);
                    if(!output.contains("ConnectionClosed"))
                    {
                        byte[] message = sut.tls.currentSeed.get(input);
                        if(message != null)
                            messageList.add(message);
                        stateList.add(targetState);
                        Word<String> inputList = curr.accessSequence.append(input);
                        Word<String> outputList = curr.outputSequence.append(output);
                        TransitionRecord newRecord = new TransitionRecord(targetState,inputList,outputList,messageList,stateList);
                        standardValue.add(messageList);
                        standardKey.add(inputList);
                        standardStatus.add(false);
                        standardScore.add(1.0);
                        total_num++;
                        seedSave(messageList, this.seed_dir + File.separator +String.format("%06d",total_num));
                        j++;
                        bfsQueue.add(newRecord);
                    }
                    else
                    {
                        byte[] message = sut.tls.currentSeed.get(input);
                        if(message != null)
                            messageList.add(message);
                        stateList.add(targetState);
                        Word<String> inputList = curr.accessSequence.append(input);
                        Word<String> outputList = curr.outputSequence.append(output);
                        standardValue.add(messageList);
                        standardKey.add(inputList);
                        standardStatus.add(true);
                        standardScore.add(1.0);
                        total_num++;
                        seedSave(messageList, this.seed_dir + File.separator +String.format("%06d",total_num));
                        j++;
                    }
                }
            }
            curr = bfsQueue.poll();
        }
        bfsQueue.clear();
        bfsQueue = null;
        String resultStr = new String();
        File file = new File(this.server_output);
        try {
            BufferedReader bufferedReader = new BufferedReader(new FileReader(file));
            String str = null;
            while (null != (str = bufferedReader.readLine())) {
                //the output of server exists many 0x00
                String tmp = str.replaceAll("\\u0000", "");
                if(tmp.contains("Coverage"))
                {
                    try {
                        String cover = tmp.substring(9).trim();
                        initCoverage = Integer.parseInt(cover);

                    } catch (NumberFormatException e) {
                        e.printStackTrace();
                    }
                }
                resultStr = resultStr.concat(tmp + "\r\n");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        FileWriter fileWriter = new FileWriter(file, false);
        fileWriter.write("");
        fileWriter.flush();
        fileWriter.close();
    }
    //list deepCopy
    public static <T> List<T> deepCopy(List<T> src)
            throws IOException, ClassNotFoundException
    {
        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
        ObjectOutputStream out = new ObjectOutputStream(byteOut);
        out.writeObject(src);
        ByteArrayInputStream byteIn = new ByteArrayInputStream(byteOut.toByteArray());
        ObjectInputStream in = new ObjectInputStream(byteIn);
        return (List<T>)in.readObject();
    }

    private static Integer readProcessOutput(InputStream inputStream, String filename) throws FileNotFoundException{
        PrintStream log = new PrintStream(new BufferedOutputStream(new FileOutputStream(filename, true)));
        Date date= new Date();
        String timestamp = String.valueOf(date.getTime());
        int length = timestamp.length();
        Integer time = 0;
        Double[] coverage = new Double[4];
        if (length > 3) {
            time = Integer.valueOf(timestamp.substring(0,length-3));
        }
        try {
            BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, Charset.forName("GBK")));
            String line;

            if ((line = reader.readLine()) != null) {
                //System.out.println(line);
                String[] tmp = line.split(" ");
                coverage[0] = Double.parseDouble(tmp[1].substring(0,tmp[1].length()-1));
                coverage[1] = Double.parseDouble(tmp[2].substring(1));
                coverage[2] = Double.parseDouble(tmp[7].substring(0,tmp[7].length()-1));
                coverage[3] = Double.parseDouble(tmp[8].substring(1));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                inputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        log.print(time);
        log.print(" ");
        log.print(coverage[0]);
        log.print(" ");
        log.print(coverage[1]);
        log.print(" ");
        log.print(coverage[2]);
        log.print(" ");
        log.println(coverage[3]);
        log.close();
        return coverage[3].intValue();
    }
    /**
     * int到byte[] 由高位到低位
     * @param i 需要转换为byte数组的整行值。
     * @return byte数组
     */
    public static byte[] intToByteArray(int i) {
        byte[] result = new byte[4];
        result[3] = (byte)((i >> 24) & 0xFF);
        result[2] = (byte)((i >> 16) & 0xFF);
        result[1] = (byte)((i >> 8) & 0xFF);
        result[0] = (byte)(i & 0xFF);
        return result;
    }

    public static Integer getRandomState(ArrayList<Double> map) {
        NavigableMap<Double, Integer> hashMap = new TreeMap<>();
        double total = 0;
        for (int i = 0 ; i < map.size(); i++ ) {
            total += map.get(i);
            hashMap.put(total,i);
        }
        Random random =new Random();
        Double weight = 1.0*random.nextInt((int)total);
        return hashMap.ceilingEntry(weight).getValue();
    }
}
